#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
合并百度新闻爬虫结果脚本
将 res 文件夹中的所有 CSV 文件合并为一个 XLSX 文件
"""

import pandas as pd
import os
import glob
from datetime import datetime

def merge_csv_to_xlsx():
    """
    合并 res 文件夹中的所有 CSV 文件为一个 XLSX 文件
    """
    # 获取 res 文件夹路径
    res_dir = "./res"
    
    # 检查 res 文件夹是否存在
    if not os.path.exists(res_dir):
        print(f"错误：{res_dir} 文件夹不存在")
        return
    
    # 获取所有 CSV 文件
    csv_files = glob.glob(os.path.join(res_dir, "*.csv"))
    
    if not csv_files:
        print(f"错误：在 {res_dir} 文件夹中没有找到 CSV 文件")
        return
    
    print(f"找到 {len(csv_files)} 个 CSV 文件：")
    for file in csv_files:
        print(f"  - {os.path.basename(file)}")
    
    # 读取并合并所有 CSV 文件
    all_dataframes = []
    
    for csv_file in sorted(csv_files):  # 按文件名排序
        try:
            print(f"正在读取：{os.path.basename(csv_file)}")
            df = pd.read_csv(csv_file)
            
            # 添加来源文件信息（可选）
            df['Source_File'] = os.path.basename(csv_file)
            
            all_dataframes.append(df)
            print(f"  ✓ 成功读取，包含 {len(df)} 行数据")
            
        except Exception as e:
            print(f"  ✗ 读取失败：{e}")
            continue
    
    if not all_dataframes:
        print("错误：没有成功读取任何 CSV 文件")
        return
    
    # 合并所有数据框
    print("\n正在合并数据...")
    merged_df = pd.concat(all_dataframes, ignore_index=True)
    print(f"合并完成，总共 {len(merged_df)} 行数据")
    
    # 生成输出文件名（包含时间戳）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f"merged_results_{timestamp}.xlsx"
    
    # 保存为 XLSX 文件
    print(f"\n正在保存到：{output_file}")
    try:
        with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
            # 保存合并后的数据到主工作表
            merged_df.to_excel(writer, sheet_name='合并数据', index=False)
            
            # 为每个原始文件创建单独的工作表（可选）
            for i, df in enumerate(all_dataframes):
                sheet_name = f"数据源_{i+1}"
                # 移除 Source_File 列，因为这是重复信息
                df_clean = df.drop('Source_File', axis=1, errors='ignore')
                df_clean.to_excel(writer, sheet_name=sheet_name, index=False)
        
        print(f"✓ 成功保存到：{output_file}")
        print(f"文件包含 {len(merged_df)} 行数据，{len(merged_df.columns)} 列")
        
        # 显示数据统计信息
        print("\n数据统计信息：")
        print(f"- 城市数量：{merged_df['City'].nunique()}")
        print(f"- 年份范围：{merged_df['Year'].min()} - {merged_df['Year'].max()}")
        
        # 统计非空数据列
        numeric_columns = merged_df.select_dtypes(include=['number']).columns
        non_empty_counts = {}
        for col in numeric_columns:
            if col not in ['Year']:  # 排除年份列
                non_empty_count = merged_df[col].notna().sum()
                if non_empty_count > 0:
                    non_empty_counts[col] = non_empty_count
        
        print(f"- 包含数据的指标列数量：{len(non_empty_counts)}")
        
    except Exception as e:
        print(f"✗ 保存失败：{e}")
        return
    
    print(f"\n合并完成！文件已保存为：{output_file}")

def main():
    """
    主函数
    """
    print("=" * 50)
    print("百度新闻爬虫结果合并工具")
    print("=" * 50)
    
    merge_csv_to_xlsx()
    
    print("\n" + "=" * 50)
    print("处理完成")
    print("=" * 50)

if __name__ == "__main__":
    main()
